package org.dromara.web.plugins.provider.impl;

import cn.hutool.core.collection.CollectionUtil;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import okhttp3.sse.EventSources;
import org.dromara.common.core.utils.StringUtils;
import org.dromara.common.mongo.ChatMessage;
import org.dromara.common.mongo.ChatMessageMongoDao;
import org.dromara.common.mongo.Session;
import org.dromara.common.mongo.SessionMongoDao;
import org.dromara.dao.business.domain.Model;
import org.dromara.web.controller.app.req.ChatBody;
import org.dromara.web.plugins.provider.ChatCompletionListener;
import org.dromara.web.plugins.provider.ChatProvider;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.List;


@Slf4j
@Component
public class ChatProviderClaude implements ChatProvider {

    public static final String MODEL_TYPE = "claude";

    private static final String COMPLETION_URL = "https://api.anthropic.com/v1/messages";
    private static final String API_KEY = "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";
    private static final String PROMPT_FM = "\\n\\nHuman: %s\\n\\nAssistant:";

    @Autowired
    OkHttpClient okHttpClient;
    @Autowired
    ThreadPoolTaskExecutor threadPoolTaskExecutor;
    @Autowired
    SessionMongoDao sessionMongoDao;
    @Autowired
    ChatMessageMongoDao chatMessageMongoDao;

    @Override
    public void chat(ChatBody body, String model, SseEmitter emitter, ChatCompletionListener listener, boolean mute) {

        EventSource.Factory eventSourceFactory = EventSources.createFactory(okHttpClient);

        JSONObject jsonObject = new JSONObject();
        if (mute) {
            jsonObject.put("stream", true);
            jsonObject.put("model", model);
            JSONArray jsonArray = new JSONArray();
            JSONObject messageObject = new JSONObject();
            messageObject.put("role", "user");
            messageObject.put("content", body.getMessage());
            jsonArray.add(messageObject);
            jsonObject.put("messages", jsonArray);
        } else {
            jsonObject = buildMessageBody(model, body.getSessionId(), body.getMessage(), body.getPrompt());
        }


        RequestBody requestBody = RequestBody.create(MediaType.parse("application/json; charset=utf-8"), jsonObject.toJSONString());

        Request request = new Request.Builder()
            .url(COMPLETION_URL)
            .header("x-api-key", API_KEY)
            .header("anthropic-version", "2023-06-01")
            .post(requestBody)
            .build();

        final StringBuilder responseBuilder = new StringBuilder();

        eventSourceFactory.newEventSource(request, new EventSourceListener() {
            @Override
            public void onEvent(@NotNull EventSource eventSource, @Nullable String id, @Nullable String type, @NotNull String data) {
                super.onEvent(eventSource, id, type, data);
                try{
                    log.info("event: {}", data);
                    emitter.send(data);
                    threadPoolTaskExecutor.execute(()->{
                        if (data.contains("[DONE]")) {

                        } else if (data.contains("choices")) {
                            log.info(data);
                            JSONObject bodyObject = JSONObject.parseObject(data);
                            JSONArray jsonArray = bodyObject.getJSONArray("choices");
                            if (jsonArray !=null && jsonArray.size() > 0) {
                                JSONObject deltaObject = jsonArray.getJSONObject(0).getJSONObject("delta");
                                if (deltaObject != null) {
                                    String content  = deltaObject.getString("content");
                                    responseBuilder.append(content == null ? "" : content);
                                }
                            }
                        }
                    });
                }catch (Exception e){

                }
            }

            @Override
            public void onClosed(@NotNull EventSource eventSource) {
                super.onClosed(eventSource);
                log.debug("closeed ....");
                emitter.complete();
                // 处理完成
                listener.onCompletion(responseBuilder.toString(), "assistant");
            }
        });
    }


    public JSONObject buildMessageBody(String model, String sessionId, String message, String prompt) {
        JSONObject jsonObject = new JSONObject();
        jsonObject.put("stream", true);
        jsonObject.put("model", model);
        JSONArray jsonArray = new JSONArray();
        if (StringUtils.isNotEmpty(prompt)) {
            jsonObject.put("prompt", prompt);
        }

        Session session = sessionMongoDao.findSession(sessionId);
        if (session != null) {
            List<ChatMessage> messageList = chatMessageMongoDao.getMessages(session.getUserId(), sessionId);
            if (CollectionUtil.isNotEmpty(messageList)) {
                messageList.forEach((chatMessage) -> {

                    JSONObject messageBody = new JSONObject();
                    messageBody.put("role", chatMessage.getRole());
                    if (chatMessage.getAsk() && StringUtils.isNotEmpty(chatMessage.getImage())) {
                        JSONArray contentArray = new JSONArray();
                        JSONObject textContent = new JSONObject();
                        textContent.put("type", "text");
                        textContent.put("text", chatMessage.getContent());
                        contentArray.add(textContent);
                        JSONObject imageContent = new JSONObject();
                        imageContent.put("type", "base64");
                        int aIndex = chatMessage.getImage().indexOf(":");
                        int bIndex = chatMessage.getImage().indexOf(";");
                        int cIndex = chatMessage.getImage().indexOf(",");
                        imageContent.put("media_type", chatMessage.getImage().substring(aIndex + 1, bIndex));
                        imageContent.put("data", chatMessage.getImage().substring(cIndex + 1));

                        contentArray.add(imageContent);
                        messageBody.put("content", contentArray);
                    } else {
                        messageBody.put("content", chatMessage.getContent());
                    }


                    jsonArray.add(messageBody);



                });
            }
        }

        jsonObject.put("messages", jsonArray);
        return jsonObject;
    }


    @Override
    public String chatResponse(Model model, String question) {
        JSONObject jsonObject = new JSONObject();

        jsonObject.put("stream", false);
        jsonObject.put("model", model.getValue());
        JSONArray jsonArray = new JSONArray();
        JSONObject messageObject = new JSONObject();
        messageObject.put("role", "user");
        messageObject.put("content", question);
        jsonArray.add(messageObject);
        jsonObject.put("messages", jsonArray);

        RequestBody requestBody = RequestBody.create(MediaType.parse("application/json; charset=utf-8"), jsonObject.toJSONString());

        Request request = new Request.Builder()
            .url(COMPLETION_URL)
            .header("Authorization", "Bearer " + API_KEY)
            .post(requestBody)
            .build();

        try {
            Response response = okHttpClient.newCall(request).execute();
            if (!response.isSuccessful()) {
                return "";
            }
            String body = response.body().string();
            JSONObject responseObject = JSONObject.parseObject(body);
            JSONArray choicesArray = responseObject.getJSONArray("choices");
            if (choicesArray != null && choicesArray.size() > 0) {
                JSONObject choiceObject = choicesArray.getJSONObject(0);
                JSONObject contentObject = choiceObject.getJSONObject("message");
                if (contentObject != null) {
                    return contentObject.getString("content");
                }
            }
        }catch (Exception e) {
            e.printStackTrace();
        }

        return "";
    }
}
